"""
主力资金建仓分析模块
分析特征：
1. 成交量放大
2. 成交额增加
3. 大单买入占比
4. 换手率变化
5. 价格走势特征
"""
from datetime import datetime, timedelta
import pandas as pd
import numpy as np
from sqlalchemy import select, and_
from models.stock_models import StockDailyData, StockRealtimeData
from database.db_engine import SessionLocal
import json

def analyze_volume_pattern(daily_data_df):
    """分析成交量模式"""
    # 计算近期成交量相对于前期的变化
    daily_data_df['volume_ma5'] = daily_data_df['volume'].rolling(window=5).mean()
    daily_data_df['volume_ma10'] = daily_data_df['volume'].rolling(window=10).mean()
    daily_data_df['volume_ratio'] = daily_data_df['volume'] / daily_data_df['volume_ma5']
    
    # 判断成交量是否放大
    volume_increase = daily_data_df['volume_ratio'] > 1.5  # 成交量是5日均量的1.5倍以上
    return volume_increase

def analyze_price_pattern(daily_data_df):
    """分析价格走势模式"""
    # 计算价格趋势指标
    daily_data_df['ma5'] = daily_data_df['close'].rolling(window=5).mean()
    daily_data_df['ma10'] = daily_data_df['close'].rolling(window=10).mean()
    daily_data_df['ma20'] = daily_data_df['close'].rolling(window=20).mean()
    
    # 计算价格波动率
    daily_data_df['price_volatility'] = (daily_data_df['high'] - daily_data_df['low']) / daily_data_df['low']
    
    # 判断是否处于横盘震荡或缓慢上涨
    price_pattern = (
        (daily_data_df['ma5'] > daily_data_df['ma20']) &  # 短期均线在长期均线上方
        (daily_data_df['price_volatility'] < 0.05)  # 日内波动率较小
    )
    return price_pattern

def analyze_big_orders(realtime_data):
    """分析大单交易情况"""
    try:
        # 解析买盘价格数据
        bid_prices = json.loads(realtime_data.bid_price) if realtime_data.bid_price else {}
        
        # 计算大单占比（假设前三档为大单）
        total_bid_volume = sum(float(vol) for vol in bid_prices.values())
        big_order_volume = sum(float(bid_prices[price]) for price in list(bid_prices.keys())[:3])
        
        big_order_ratio = big_order_volume / total_bid_volume if total_bid_volume > 0 else 0
        return big_order_ratio > 0.6  # 大单占比超过60%
    except:
        return False

def get_main_force_stocks(days=20, volume_threshold=1.5):
    """
    识别可能存在主力建仓的股票
    
    Args:
        days (int): 分析的历史数据天数
        volume_threshold (float): 成交量放大的阈值
    
    Returns:
        list: 疑似主力建仓的股票列表
    """
    end_date = datetime.now().date()
    start_date = end_date - timedelta(days=days)
    
    session = SessionLocal()
    try:
        # 1. 获取历史日线数据
        daily_query = select(StockDailyData).where(
            and_(
                StockDailyData.trade_date >= start_date,
                StockDailyData.trade_date <= end_date
            )
        )
        daily_results = session.execute(daily_query).scalars().all()
        
        # 2. 获取最新实时数据
        realtime_query = select(StockRealtimeData).where(
            StockRealtimeData.trade_time >= datetime.now() - timedelta(hours=4)
        )
        realtime_results = session.execute(realtime_query).scalars().all()
        
        # 转换为DataFrame
        daily_df = pd.DataFrame([{
            'ts_code': r.ts_code,
            'trade_date': r.trade_date,
            'close': r.close,
            'high': r.high,
            'low': r.low,
            'volume': r.volume,
            'amount': r.amount,
            'turnover_rate': r.turnover_rate
        } for r in daily_results])
        
        if daily_df.empty:
            return []
        
        # 按股票代码分组分析
        result_stocks = []
        for ts_code, group in daily_df.groupby('ts_code'):
            group = group.sort_values('trade_date')
            
            # 1. 分析成交量模式
            volume_signal = analyze_volume_pattern(group)
            
            # 2. 分析价格走势
            price_signal = analyze_price_pattern(group)
            
            # 3. 获取该股票的最新实时数据
            realtime_data = next((r for r in realtime_results if r.ts_code == ts_code), None)
            big_order_signal = analyze_big_orders(realtime_data) if realtime_data else False
            
            # 4. 换手率分析
            turnover_signal = group['turnover_rate'].iloc[-1] > group['turnover_rate'].mean() * 1.2
            
            # 综合判断
            if (volume_signal.iloc[-1] and  # 成交量放大
                price_signal.iloc[-1] and   # 价格形态符合
                turnover_signal and         # 换手率上升
                (big_order_signal or not realtime_data)):  # 大单买入（如果有实时数据）
                
                result_stocks.append({
                    'ts_code': ts_code,
                    'last_date': group['trade_date'].iloc[-1],
                    'volume_ratio': group['volume_ratio'].iloc[-1],
                    'turnover_rate': group['turnover_rate'].iloc[-1],
                    'price_volatility': group['price_volatility'].iloc[-1]
                })
        
        return result_stocks
    
    finally:
        session.close()

def print_analysis_results():
    """打印分析结果"""
    results = get_main_force_stocks()
    if not results:
        print("没有发现疑似主力建仓的股票")
        return
    
    print(f"\n发现 {len(results)} 只疑似主力建仓的股票：")
    print("=" * 60)
    for stock in results:
        print(f"股票代码: {stock['ts_code']}")
        print(f"最新交易日: {stock['last_date']}")
        print(f"成交量放大倍数: {stock['volume_ratio']:.2f}")
        print(f"换手率: {stock['turnover_rate']:.2f}%")
        print(f"价格波动率: {stock['price_volatility']:.2f}%")
        print("-" * 60)

if __name__ == "__main__":
    print_analysis_results() 